! module *tree_struct*
!
! Purpose:
!   (1) construct tree structure for storing information
!   (2) bitree, quadtree, octree correspond to different dimensions
!
! Author:
!   Name: DONG Li
!   Email: dongli@lasg.iap.ac.cn

module tree_struct

    use space_manager
    use data_fit, only: data_fit_min_num_data
    use message
    use run_manager

    implicit none

    type tree_node
        ! hierarchical structure
        integer :: lev
        integer :: root_index(3)
        integer :: num_child(3)
        type(tree_node), pointer :: child(:,:,:)
        logical :: end_node
        ! geometrical properties
        real, allocatable :: xc(:) ! location of center point
        real, allocatable :: r(:)  ! radius of influence ellipsoid
        ! data index structure
        integer :: np  ! total number of parcels
        integer :: npi ! number of parcels in main box
        integer :: npo ! number of parcels in influence ellipsoid
        integer, allocatable :: pi(:) ! parcel index
        ! for periodic boundary condition
        integer, allocatable :: pb(:,:)
        ! approximation coefficients 
        integer :: approx_order ! different subdomain may have different order
        real, allocatable :: c(:)
    end type tree_node

    type forest_node
        type(tree_node), allocatable :: tree(:,:,:)
    end type forest_node

    integer :: max_num_parcel
    integer :: max_approx_order

    integer :: num_tree(3)
    integer :: num_forest
    type(forest_node), allocatable, target :: forest(:)

    integer :: max_lev
    real, allocatable :: box_len(:,:)

    real :: ratio ! box_len/2/r
    real :: rate  ! incremental rate of r

    namelist /tree_struct_parameters/ &
        max_num_parcel, num_tree, max_lev, ratio, rate

contains

    subroutine tree_struct_init(namelist_file, num_forest_, approx_order)
        character(*), intent(in) :: namelist_file
        integer, intent(in) :: num_forest_
        integer, intent(in) :: approx_order

        integer :: i, j, k, l, m, o(3)

        ! default values
        max_num_parcel = 100
        max_lev = 10
        ratio = 0.9
        rate  = 0.1
        
        ! read in parameters from namelist file
        open(10, file=namelist_file)
        read(10, nml=tree_struct_parameters, end=10)
10      close(10)

        do i = 1, num_dim
            if(num_tree(i) == 0) then
                call message_show(ERROR_MSG, &
                    "tree_struct",           &
                    "tree_struct_init",      &
                    "invalid num_tree in namelist")
                call run_manager_end_run
            end if
        end do

        num_forest = num_forest_
        max_approx_order = approx_order

        allocate(forest(num_forest))
        do i = 1, num_forest
            allocate(forest(i)%tree(num_tree(1),num_tree(2),num_tree(3)))
        end do
        allocate(box_len(num_dim,max_lev))

        ! set the tree box lengthes
        do i = 1, num_dim
            box_len(i,1) = (boundary(2,i)-boundary(1,i))/num_tree(i)
        end do
        ! 
        do j = 2, max_lev
            box_len(:,j) = box_len(:,j-1)/2.0
        end do

        ! initialize tree structure for each forest
        do m = 1, num_forest
        do k = 1, num_tree(3)
        do j = 1, num_tree(2)
        do i = 1, num_tree(1)
            o = (/i,j,k/)
            forest(m)%tree(i,j,k)%lev = 1

            forest(m)%tree(i,j,k)%root_index = (/i,j,k/)
            forest(m)%tree(i,j,k)%num_child = 0
            forest(m)%tree(i,j,k)%end_node = .true.

            allocate(forest(m)%tree(i,j,k)%xc(num_dim))
            allocate(forest(m)%tree(i,j,k)%r(num_dim))
            do l = 1, num_dim
                forest(m)%tree(i,j,k)%xc(l) = boundary(1,l)+(o(l)-0.5)*box_len(l,1)
            end do
            forest(m)%tree(i,j,k)%r = box_len(:,1)/2/ratio

            forest(m)%tree(i,j,k)%npi = 0
            forest(m)%tree(i,j,k)%npo = 0
            allocate(forest(m)%tree(i,j,k)%pi(max_num_parcel))
            allocate(forest(m)%tree(i,j,k)%pb(num_dim,max_num_parcel))
            forest(m)%tree(i,j,k)%pb = 0

            forest(m)%tree(i,j,k)%approx_order = approx_order
        end do
        end do
        end do
        end do

        return
    end subroutine tree_struct_init

    subroutine tree_struct_finalize
        
    end subroutine tree_struct_finalize

    subroutine tree_struct_reset
        integer :: i, j, k, l

        do l = 1, num_forest
        do k = 1, num_tree(3)
        do j = 1, num_tree(2)
        do i = 1, num_tree(1)
            forest(l)%tree(i,j,k)%np = 0
            forest(l)%tree(i,j,k)%npi = 0
            forest(l)%tree(i,j,k)%npo = 0
            forest(l)%tree(i,j,k)%end_node = .true.
            forest(l)%tree(i,j,k)%r = box_len(:,1)/2/ratio
            forest(l)%tree(i,j,k)%approx_order = max_approx_order
        end do
        end do
        end do
        end do

    end subroutine tree_struct_reset

    subroutine tree_struct_where_point(x, o)
        real, intent(in) :: x(num_dim)
        integer, intent(out) :: o(3)

        integer :: i

        o = 1
        do i = 1, num_dim
            o(i) = ceiling(x(i)/box_len(i,1))
            if(x(i) == boundary(1,i)) o(i) = 1
            if(x(i) == boundary(2,i)) o(i) = num_tree(i)
            if(o(i) > num_tree(i) .or. o(i) < 1) then
                call message_show(ERROR_MSG,     &
                    "tree_struct",               &
                    "tree_struct_where_point", &
                    "data point location exceeds boundary")
                call run_manager_end_run
            end if
        end do

        return
    end subroutine tree_struct_where_point

    subroutine tree_struct_boundary_check(o1, o2)
        integer, intent(in) :: o1(3)
        integer, intent(out) :: o2(3)

        integer :: i

        o2 = o1
        ! boundary condition
        do i = 1, num_dim
            select case(boundary_type(i))
            case("periodic")
                if(o1(i) < 1) o2(i) = num_tree(i)+o1(i)
                if(o1(i) > num_tree(i)) o2(i) = o1(i)-num_tree(i)
            case default
                o2(i) = max(1, o1(i))
                o2(i) = min(num_tree(i), o1(i))
            end select
        end do

        return
    end subroutine tree_struct_boundary_check

    subroutine tree_struct_register_data(np, x)
        integer, intent(in) :: np
        real, intent(in) :: x(num_dim,np)

        integer :: i, j, k, l
        type(tree_node), pointer :: tree1(:,:,:)
        type(tree_node), pointer :: tree2(:,:,:)
        integer :: o(3)         ! tree node index
        integer :: n(3)         ! number of covered nodes
        integer :: ir(3)        ! reference index (one corner)
        logical :: checked(num_tree(1),num_tree(2),num_tree(3))
        integer :: counter

        ! tree in each forest is identical, so choose first one
        ! to perform operation, and copy result to others

        tree1 => forest(1)%tree

        ! put data into tree main box
        do j = 1, np
            call tree_struct_where_point(x(:,j), o)
            i = tree1(o(1),o(2),o(3))%npi+1
            tree1(o(1),o(2),o(3))%npi = i
            tree1(o(1),o(2),o(3))%pi(i) = j
        end do

        ! put outside-box data into subdomain
        do k = 1, num_tree(3)
        do j = 1, num_tree(2)
        do i = 1, num_tree(1)
            counter = 1
            do while(.true.)
                checked = .false.
                ! ignore itself
                checked(i,j,k) = .true.
                n = 1
                ir = (/i,j,k/)
                do l = 1, num_dim
                    n(l) = int((2*tree1(i,j,k)%r(l)-box_len(l,1))/box_len(l,1))+3
                    ir(l) = ir(l)-n(l)/2
                end do
                call distance_check(tree1(i,j,k), ir, n, checked, np, x)
                if(tree1(i,j,k)%npi+tree1(i,j,k)%npo < &
                    data_fit_min_num_data(tree1(i,j,k)%approx_order)) then
                    tree1(i,j,k)%r = tree1(i,j,k)%r*(1+rate)
                    counter = counter+1
                    cycle
                else
                    exit
                end if
            end do
            tree1(i,j,k)%np = tree1(i,j,k)%npi+tree1(i,j,k)%npo
        end do
        end do
        end do

        ! copy first tree to other trees
        do l = 2, num_forest
            tree2 => forest(l)%tree
            do k = 1, num_tree(3)
            do j = 1, num_tree(2)
            do i = 1, num_tree(1)
                tree2(i,j,k)%r = tree1(i,j,k)%r
                tree2(i,j,k)%np = tree1(i,j,k)%np
                tree2(i,j,k)%npi = tree1(i,j,k)%npi
                tree2(i,j,k)%npo = tree1(i,j,k)%npo
                tree2(i,j,k)%pi = tree1(i,j,k)%pi
                tree2(i,j,k)%pb = tree1(i,j,k)%pb
            end do
            end do
            end do
        end do

        return
    end subroutine tree_struct_register_data

    subroutine distance_check(node, ir, n, checked, np, x)
        ! see the comments in tree_struct_register_data
        integer, intent(in) :: ir(3), n(3), np
        type(tree_node), intent(inout) :: node
        logical, intent(inout) :: checked(num_tree(1),num_tree(2),num_tree(3))
        real, intent(in) :: x(num_dim,np)

        integer :: i, j, k, l
        integer :: o1(3), o2(3)
        integer :: i1, j1, k1
        integer :: i2, j2, k2
        real :: x0(num_dim), x1(num_dim)
        real :: temp, temp1, temp2
        integer :: dislocation(num_dim)
        logical :: flag
        ! temporal method
        type(tree_node), pointer :: tree(:,:,:)

        tree => forest(1)%tree

        x0 = node%xc
        do k1 = ir(3), ir(3)+n(3)-1
        do j1 = ir(2), ir(2)+n(2)-1
        do i1 = ir(1), ir(1)+n(1)-1
            o1 = (/i1,j1,k1/)
            call tree_struct_boundary_check(o1, o2)
            if(checked(o2(1),o2(2),o2(3))) cycle
            ! check the data points contained in the main box
            do k = 1, tree(o2(1),o2(2),o2(3))%npi
                ! check whether the point is included in subdomain or not
                flag = .false.
                l = tree(o2(1),o2(2),o2(3))%pi(k)
                do j = node%npi+1, node%npi+node%npo
                    if(l == node%pi(j)) then
                        flag = .true.
                        exit
                    end if
                end do
                if(flag == .true.) cycle
                x1 = x(:,l)
                dislocation = 0
                ! ellipsoid equation
                temp = 0.0
                do j = 1, num_dim
                    temp1 = abs(x1(j)-x0(j))
                    select case(boundary_type(j))
                    case("periodic")
                        temp2 = boundary(2,j)-boundary(1,j)-temp1
                        if(temp1 > temp2) then
                            temp = temp+(temp2/node%r(j))**2
                            if(x1(j) > x0(j)) then
                                dislocation(j) = -1
                            else
                                dislocation(j) = 1
                            end if
                        else
                            temp = temp+(temp1/node%r(j))**2
                        end if
                    case default
                        temp = temp+(temp1/node%r(j))**2
                    end select
                end do
                ! in ellipsoid
                if(temp <= 1) then
                    node%npo = node%npo+1
                    j = node%npi+node%npo
                    node%pi(j) = l
                    node%pb(:,j) = dislocation
                end if
            end do
            checked(o2(1),o2(2),o2(3)) = .true.
        end do
        end do
        end do

        return
    end subroutine distance_check

    subroutine tree_struct_subdivide_tree(node, np, x)
        type(tree_node), target, intent(inout) :: node
        integer, intent(in) :: np
        real, intent(in) :: x(np)

        real :: xr(num_dim)             ! reference location
        integer :: ir(3)                ! reference index
        integer :: i, j, k, l
        integer :: o(3), n(3)
        integer :: lev, counter
        logical :: checked(num_tree(1),num_tree(2),num_tree(3))

        node%end_node = .false.
        
        lev = node%lev+1

        ! the left and lower corner will be the reference point
        xr = node%xc-box_len(:,lev-1)/2.0
 
        ! clean up unused data
        deallocate(node%c)

        if(associated(node%child) == .false.) then
            select case(num_dim)
            case(1)
                node%num_child = (/2,1,1/) 
            case(2)
                node%num_child = (/2,2,1/) 
            case(3)
                node%num_child = (/2,2,2/) 
            case default
                ! handle error
            end select
            allocate(node%child(node%num_child(1),node%num_child(2),node%num_child(3)))

            ! initialize childs
            do k = 1, node%num_child(3)
            do j = 1, node%num_child(2)
            do i = 1, node%num_child(1)
                o = (/i,j,k/)
                node%child(i,j,k)%lev = lev
 
                node%child(i,j,k)%root_index = node%root_index
                node%child(i,j,k)%num_child  = 0
                node%child(i,j,k)%end_node   = .true.
 
                allocate(node%child(i,j,k)%xc(num_dim))
                allocate(node%child(i,j,k)%r(num_dim))
                do l = 1, num_dim
                    node%child(i,j,k)%xc(l) = xr(l)+(o(l)-0.5)*box_len(l,lev)
                end do
                node%child(i,j,k)%r = box_len(:,lev)/2/ratio
 
                node%child(i,j,k)%npi = 0
                node%child(i,j,k)%npo = 0
                allocate(node%child(i,j,k)%pi(max_num_parcel))
                allocate(node%child(i,j,k)%pb(num_dim,max_num_parcel))
                node%child(i,j,k)%pb = 0

                node%child(i,j,k)%approx_order = max_approx_order
            end do
            end do
            end do
        else
            do k = 1, node%num_child(3)
            do j = 1, node%num_child(2)
            do i = 1, node%num_child(1)
                node%child(i,j,k)%end_node = .true.
                node%child(i,j,k)%npi = 0
                node%child(i,j,k)%npo = 0
            end do
            end do
            end do
        end if
 
        do k = 1, node%num_child(3)
        do j = 1, node%num_child(2)
        do i = 1, node%num_child(1)
            counter = 1
            ! --> needs optimization
            do while(.true.)
                checked = .false.
                n = 1
                ir = node%root_index
                do l = 1, num_dim
                    n(l) = int((2*node%child(i,j,k)%r(l)-box_len(l,lev))/box_len(l,1))+3
                    ir(l) = ir(l)-n(l)/2
                end do
                call distance_check(node%child(i,j,k), ir, n, checked, np, x)
                if(node%child(i,j,k)%npo < &
                    data_fit_min_num_data(node%child(i,j,k)%approx_order)) then
                    node%child(i,j,k)%r = node%child(i,j,k)%r*(1+rate)
                    counter = counter+1
                    cycle
                else
                    exit
                end if
            end do
            node%child(i,j,k)%np = node%child(i,j,k)%npo
        end do
        end do
        end do

        return
    end subroutine tree_struct_subdivide_tree

end module tree_struct
